package com.iraytek.rtsplib;

import android.content.Context;
import android.media.MediaCodec;
import android.media.MediaFormat;
import android.media.MediaMuxer;
import android.os.Handler;
import android.os.HandlerThread;

import androidx.annotation.IntRange;

import java.io.IOException;
import java.nio.ByteBuffer;

import pedro.encoder.BaseEncoder;
import pedro.encoder.audio.AudioEncoder;
import pedro.encoder.audio.GetAacData;
import pedro.encoder.input.audio.MicrophoneManager;
import pedro.encoder.input.audio.MicrophoneManagerManual;
import pedro.encoder.input.video.CameraHelper;
import pedro.encoder.utils.CodecUtil;
import pedro.encoder.video.FormatVideoEncoder;
import pedro.encoder.video.GetVideoData;
import pedro.encoder.video.VideoEncoder;

/**
 * used for video/audio encode management
 */
public abstract class BaseEncodeManager implements GetVideoData, GetAacData {
    private static final int PREP_FLAG_VIDEO = 1;
    private static final int PREP_FLAG_AUDIO = 2;
    private int prepareFlag = 0;
    private Context context;
    protected VideoEncoder videoEncoder;
    private AudioEncoder audioEncoder;
    private MicrophoneManager microphoneManager;
    private boolean streaming;
    private FormatVideoEncoder videoEncodeFormat = FormatVideoEncoder.SURFACE;

    protected boolean recording;
    private int videoTrackIndex = -1, audioTrackIndex = -1;
    private MediaMuxer mediaMuxer;
    private MediaFormat videoFormat, audioFormat;
    private Handler handler;

    public BaseEncodeManager(Context context) {
        this.context = context;
        init();
    }

    private void init() {
        videoEncoder = createVideoEncoder(this);
        audioEncoder = new AudioEncoder(this);
        microphoneManager = new MicrophoneManagerManual();//or use MicrophoneManager
        audioEncoder.setGetFrame(((MicrophoneManagerManual) microphoneManager).getGetFrame());
    }

    /**
     * Set hardware encoding or software encoding, must be called before
     * must be called before {@link #prepareVideo()} & {@link #prepareAudio()}
     *
     * @param force encoding mode, default is {@link CodecUtil.Force#FIRST_COMPATIBLE_FOUND}
     */
    public void setForce(CodecUtil.Force force) {
        if (videoEncoder != null) {
            videoEncoder.setForce(force);
        }
        if (audioEncoder != null) {
            audioEncoder.setForce(force);
        }
    }

    public boolean isStreaming() {
        return streaming;
    }

    /**
     * Use default params for vudio.
     * Call this method before use{@link #startStream(String)}
     *
     * @return true if success, false if you get a error
     */
    public boolean prepareVideo() {
        int rotation = CameraHelper.getCameraOrientation(context);
        return prepareVideo(1280, 720, 30, 1200 * 1024, rotation);
    }

    public boolean prepareVideo(int width, int height, @IntRange(from = 1, to = 32) int quality) {
        int rotation = CameraHelper.getCameraOrientation(context);
        return prepareVideo(width, height, 30, width * height * quality, rotation);
    }

    public boolean prepareVideo(int width, int height, int fps, int bitrate, int rotation) {
        return prepareVideo(width, height, fps, bitrate, 2, rotation);
    }

    public boolean prepareVideo(int width, int height, int fps, int bitrate, int iFrameInterval,
                                int rotation) {
        return prepareVideo(width, height, fps, bitrate, iFrameInterval, rotation, -1, -1);
    }

    /**
     * Call this method before use @startStream. If not you will do a stream without video.
     *
     * @param width    resolution in px.
     * @param height   resolution in px.
     * @param fps      frames per second of the stream.
     * @param bitrate  H264 in bps.
     * @param rotation could be 90, 180, 270 or 0 (Normally 0 if you are streaming in landscape or 90
     *                 if you are streaming in Portrait). This only affect to stream result. NOTE: Rotation with
     *                 encoder is silence ignored in some devices.
     * @return true if success, false if you get a error (Normally because the encoder selected
     * doesn't support any configuration seated or your device hasn't a H264 encoder).
     */
    public boolean prepareVideo(int width, int height, int fps, int bitrate, int iFrameInterval,
                                int rotation, int avcProfile, int avcProfileLevel) {

        boolean success = videoEncoder.prepareVideoEncoder(width, height, fps, bitrate, rotation, iFrameInterval,
                videoEncodeFormat, avcProfile, avcProfileLevel);
        if (success) {
            prepareFlag |= PREP_FLAG_VIDEO;
        }
        return success;
    }

    /**
     * must be called before{@link #prepareVideo()}
     *
     * @param videoEncodeFormat see{@link FormatVideoEncoder}
     */
    public BaseEncodeManager setVideoEncodeFormat(FormatVideoEncoder videoEncodeFormat) {
        this.videoEncodeFormat = videoEncodeFormat;
        return this;
    }

    /**
     * set size for encoding, this will set width & height and restart video-encoder
     */
    public void setVideoEncodeSize(int width, int height) {
        videoEncoder.setSize(width, height);
        if (isStreaming()) {
            videoEncoder.reset();
        }
    }

    /**
     * @param type video format,such as h264,h265 etc. see {@link pedro.encoder.utils.CodecUtil}
     */
    public void setVideoFormat(String type) {
        videoEncoder.setType(type);
    }

    /**
     * Use default params for audio.
     * Call this method before use{@link #startStream(String)}
     *
     * @return true if success, false if you get a error
     */
    public boolean prepareAudio() {
        return prepareAudio(64 * 1024, 32000, true, false, false);
    }

    /**
     * Call this method before use @startStream. If not you will do a stream without audio.
     *
     * @param bitrate         AAC in kb.
     * @param sampleRate      of audio in hz. Can be 8000, 16000, 22500, 32000, 44100.
     * @param isStereo        true if you want Stereo audio (2 audio channels), false if you want Mono audio
     *                        (1 audio channel).
     * @param echoCanceler    true enable echo canceler, false disable.
     * @param noiseSuppressor true enable noise suppressor, false  disable.
     * @return true if success, false if you get a error (Normally because the encoder selected
     * doesn't support any configuration seated or your device hasn't a AAC encoder).
     */
    public boolean prepareAudio(int bitrate, int sampleRate, boolean isStereo, boolean echoCanceler,
                                boolean noiseSuppressor) {
        if (!microphoneManager.createMicrophone(sampleRate, isStereo, echoCanceler, noiseSuppressor)) {
            return false;
        }
        prepareAudioRtp(isStereo, sampleRate);
        boolean success = audioEncoder.prepareAudioEncoder(bitrate, sampleRate, isStereo,
                microphoneManager.getMaxInputSize());
        if (success) {
            prepareFlag |= PREP_FLAG_AUDIO;
        }
        return success;
    }

    /**
     * Record video&audio with format MPEG-4, must be called after video&audio prepared
     *
     * @param path path to save video
     */
    public void startRecording(String path) {
        if (recording) {
            return;
        }
        if (videoEncoder.getStatus() == BaseEncoder.Status.PREPARED) {
            startEncoders();
        } else if (!streaming) {
            resetEncoders();
        }
        try {
            mediaMuxer = new MediaMuxer(path, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4);
            recording = true;
            HandlerThread thread = new HandlerThread("record-prepare");
            thread.start();
            handler = new Handler(thread.getLooper());

            addTracks(mediaMuxer);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private void addTracks(MediaMuxer mediaMuxer) {
        if ((prepareFlag & PREP_FLAG_VIDEO) == PREP_FLAG_VIDEO) {
            //make sure videoFormat added to track
            if (videoTrackIndex == -1) {
                if (videoFormat != null) {
                    videoTrackIndex = mediaMuxer.addTrack(videoFormat);
                } else {
                    handler.postDelayed(() -> addTracks(mediaMuxer), 50);
                    return;
                }
            }
        }

        if ((prepareFlag & PREP_FLAG_AUDIO) == PREP_FLAG_AUDIO) {
            //make sure audioFormat added to track
            if (audioTrackIndex == -1) {
                if (audioFormat != null) {
                    audioTrackIndex = mediaMuxer.addTrack(audioFormat);
                } else {
                    handler.postDelayed(() -> addTracks(mediaMuxer), 50);
                    return;
                }
            }
        }
        mediaMuxer.start();
    }

    /**
     * Stop video&audio recording, and also stop encoders if it's not streaming
     */
    public void stopRecording() {
        handler.getLooper().quit();
        if (recording) {
            recording = false;
            videoFormat = null;
            audioFormat = null;
            videoTrackIndex = audioTrackIndex = -1;
            if (mediaMuxer != null) {
                try {
                    mediaMuxer.stop();
                    mediaMuxer.release();
                } catch (Exception e) {
                    e.printStackTrace();
                }
                mediaMuxer = null;
            }
        }
        if (!streaming) {
            stopEncoders();
        }
    }

    /**
     * start encoder, need to be called after {@link #prepareVideo()} or/and {@link #prepareAudio()}
     */
    public void startStream() {
        startStream("");
    }

    /**
     * start encoder, need to be called after {@link #prepareVideo()} or/and {@link #prepareAudio()}
     *
     * @param url of the stream like: protocol://ip:port/application/streamName
     */
    public void startStream(String url) {
        if (!streaming) {
            streaming = true;
            if (videoEncoder.getStatus() == BaseEncoder.Status.PREPARED) {
                startEncoders();
            } else if (!recording) {
                resetEncoders();
            }
            startStreamRtp(url);
        }
    }

    public void stopStream() {
        if (streaming) {
            streaming = false;
            stopStreamRtp();
        }
        if (!recording) {
            stopEncoders();
        }
    }


    private void startEncoders() {
        if ((prepareFlag & PREP_FLAG_VIDEO) == PREP_FLAG_VIDEO) {
            if (!videoEncoder.isRunning()) {
                videoEncoder.start();
            }
        }
        if ((prepareFlag & PREP_FLAG_AUDIO) == PREP_FLAG_AUDIO) {
            if (!audioEncoder.isRunning()) {
                audioEncoder.start();
                microphoneManager.start();
            }
        }
    }

    private void resetEncoders() {
        if ((prepareFlag & PREP_FLAG_VIDEO) == PREP_FLAG_VIDEO) {
            videoEncoder.reset();
        }
        if ((prepareFlag & PREP_FLAG_AUDIO) == PREP_FLAG_AUDIO) {
            audioEncoder.reset();
        }
    }

    private void stopEncoders() {
        videoEncoder.stop();
        audioEncoder.stop();
        microphoneManager.stop();
    }

    protected String getRtspAddress() {
        return null;
    }

    /**
     * set params of audio stream
     * 设置音频流参数
     *
     * @param isStereo   立体声
     * @param sampleRate 采样率
     */
    protected abstract void prepareAudioRtp(boolean isStereo, int sampleRate);

    /**
     * encoded audio data come, for next step, send audio stream
     */
    protected abstract void getAacDataRtp(ByteBuffer aacBuffer, MediaCodec.BufferInfo info);

    /**
     * format change, usually on start of stream, invoked from callback in {@link MediaCodec#setCallback(MediaCodec.Callback)}
     *
     * @param sps 序列参数集
     * @param pps 图像参数集
     * @param vps 视频参数集 for h265
     */
    protected abstract void onSpsPpsVpsRtp(ByteBuffer sps, ByteBuffer pps, ByteBuffer vps);

    /**
     * encoded video data come
     */
    protected abstract void getVideoDataRtp(ByteBuffer buffer, MediaCodec.BufferInfo info);

    /**
     * override this to use custom videoEncoder
     */
    protected VideoEncoder createVideoEncoder(GetVideoData getVideoData) {
        return new VideoEncoder(getVideoData);
    }

    /**
     * generally used for PUSH stream
     *
     * @param url of the stream like: protocol://ip:port/application/streamName
     */
    protected void startStreamRtp(String url) {
    }

    protected abstract void stopStreamRtp();

    @Override
    public void onSpsPps(ByteBuffer sps, ByteBuffer pps) {
        if (streaming) onSpsPpsVpsRtp(sps, pps, null);
    }

    @Override
    public void onSpsPpsVps(ByteBuffer sps, ByteBuffer pps, ByteBuffer vps) {
        if (streaming) onSpsPpsVpsRtp(sps, pps, vps);
    }


    @Override
    public void getVideoData(ByteBuffer buffer, MediaCodec.BufferInfo info) {
        if (streaming) getVideoDataRtp(buffer, info);
        if (recording) {
            if (videoTrackIndex != -1) {
                mediaMuxer.writeSampleData(videoTrackIndex, buffer, info);
            }
        }
    }


    @Override
    public void onVideoFormat(MediaFormat mediaFormat) {
        videoFormat = mediaFormat;
    }

    //<------------------------audio-----------------------------------------------
    @Override
    public void getAacData(ByteBuffer aacBuffer, MediaCodec.BufferInfo info) {
        if (streaming) getAacDataRtp(aacBuffer, info);
        if (recording) {
            if (audioTrackIndex != -1) {
                mediaMuxer.writeSampleData(audioTrackIndex, aacBuffer, info);
            }
        }
    }


    @Override
    public void onAudioFormat(MediaFormat mediaFormat) {
        audioFormat = mediaFormat;
    }
    //------------------------audio----------------------------------------------->
}
